'''
Author: kkk 997610780@qq.com
Date: 2025-02-28 16:13:39
LastEditors: kkk 997610780@qq.com
LastEditTime: 2025-03-06 14:29:40
FilePath: \get-video-images\app.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''

from flask import Flask, request, jsonify, send_from_directory
import os
import uuid  # 确保这行导入存在
import cv2
from werkzeug.utils import secure_filename
from videoAction import extract_frames,download_video,get_video_url_from_share
from pathlib import Path

app = Flask(__name__)

# 配置上传文件存储路径
# UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads')
# if not os.path.exists(UPLOAD_FOLDER):
#     os.makedirs(UPLOAD_FOLDER)

    
try:
    # 尝试创建网络共享文件夹路径
    UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads')
    if not os.path.exists(UPLOAD_FOLDER):
        os.makedirs(UPLOAD_FOLDER)
except Exception as e:
    print(f"无法访问网络共享文件夹: {str(e)}")
    # 使用本地备用路径
    UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads')
    if not os.path.exists(UPLOAD_FOLDER):
        os.makedirs(UPLOAD_FOLDER)
    print(f"已切换到本地备用路径: {UPLOAD_FOLDER}")

# 允许的文件扩展名
ALLOWED_EXTENSIONS = {'mp4', 'avi', 'mov', 'wmv', 'flv', 'mkv'}

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

# 添加上传图片接口

# 在 ALLOWED_EXTENSIONS 定义后添加图片允许的扩展名
ALLOWED_IMAGE_EXTENSIONS = {'jpg', 'jpeg', 'png', 'gif', 'bmp', 'webp'}

# 在 allowed_file 函数后添加新的函数
def allowed_image(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_IMAGE_EXTENSIONS

# 添加新的上传图片接口
@app.route('/upload/image', methods=['POST'])
def upload_image():
    """
    上传图片接口：接受图片文件并保存
    """
    try:
        # 检查是否有文件上传
        if 'file' in request.files:
            file = request.files['file']
            
            # 如果用户没有选择文件，浏览器也会提交一个空的文件部分
            if file.filename == '':
                return jsonify({'error': '没有选择文件'}), 400
            
            if file and allowed_image(file.filename):
                # 安全地获取文件名并保存文件
                filename = secure_filename(file.filename)
                # 添加唯一标识符，避免文件名冲突
                unique_filename = f"{uuid.uuid4().hex}_{filename}"
                file_path = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
                file.save(file_path)
                
                return jsonify({
                    'success': True,
                    'message': '图片上传成功',
                    'file_path': file_path,
                    'filename': unique_filename
                })
            
            return jsonify({'error': '不允许的文件类型'}), 400
        
        else:
            return jsonify({'error': '请提供图片文件'}), 400
            
    except Exception as e:
        return jsonify({'error': f'上传图片时出错: {str(e)}'}), 500

@app.route('/upload', methods=['POST'])
def upload_video():
    """
    上传视频接口：接受视频文件上传并处理
    """
    try:
        # 检查是否有文件上传
        if 'file' in request.files:
            file = request.files['file']
            
            # 如果用户没有选择文件，浏览器也会提交一个空的文件部分
            if file.filename == '':
                return jsonify({'error': '没有选择文件'}), 400
            
            if file and allowed_file(file.filename):
                # 安全地获取文件名并保存文件
                filename = secure_filename(file.filename)
                # 添加唯一标识符，避file_path免文件名冲突
                unique_filename = f"{uuid.uuid4().hex}_{filename}"
                file_path = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
                
                file.save(file_path)
                
                # 处理视频，提取帧
                output_dir = extract_frames(file_path)
                
                # 处理完成后删除上传的视频文件
                if os.path.exists(file_path):
                    os.remove(file_path)
                
                return jsonify({
                    'success': True,
                    'message': '视频处理成功',
                    'output_directory': output_dir
                })
            
            return jsonify({'error': '不允许的文件类型'}), 400
        
        else:
            return jsonify({'error': '请提供视频文件'}), 400
            
    except Exception as e:
        return jsonify({'error': f'处理视频时出错: {str(e)}'}), 500

@app.route('/frames/<path:folder_name>/<path:filename>', methods=['GET'])
def get_frame(folder_name, filename):
    """获取处理后的视频帧图片"""
    frames_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "frames")
    return send_from_directory(os.path.join(frames_dir, folder_name), filename)

# 新增访问上传文件夹视频的路由
@app.route('/uploads/<path:filename>', methods=['GET'])
def get_video(filename):
    """获取上传文件夹中的视频文件"""
    try:
        return send_from_directory(app.config['UPLOAD_FOLDER'], filename)
    except Exception as e:
        return jsonify({'error': f'获取视频文件失败: {str(e)}'}), 404

@app.route('/health', methods=['GET'])
def health_check():
    return jsonify({'status': 'ok'})

@app.route('/download', methods=['POST'])
def download_from_url():
    """
    通过抖音视频链接下载视频并处理
    """
    try:
        # 检查Content-Type
        if not request.is_json:
            return jsonify({
                'error': '请求格式错误',
                'detail': f'当前Content-Type: {request.content_type}, 需要: application/json'
            }), 415
            
        data = request.get_json(silent=True)
        if data is None:
            return jsonify({
                'error': '无效的JSON数据',
                'detail': '请确保请求体是有效的JSON格式'
            }), 400
            
        if 'url' not in data:
            return jsonify({'error': '请提供视频链接'}), 400
            
        video_url = data['url']
        
        # 如果是分享链接，先获取实际视频地址
        if 'v.douyin.com' in video_url or 'www.douyin.com' in video_url:
            real_video_url = get_video_url_from_share(video_url)
            if not real_video_url:
                return jsonify({'error': '无法获取视频地址'}), 400
            video_url = real_video_url
        
        # 下载视频
        video_filename = f"{uuid.uuid4().hex}.mp4"
        video_path = os.path.join(app.config['UPLOAD_FOLDER'], video_filename)
        print('视频地址',video_path)
        if not download_video(video_url, video_path):
            return jsonify({'error': '视频下载失败'}), 400
            
        # 处理视频，提取帧
        output_dir = extract_frames(video_path)
        
        # 处理完成后删除临时视频文件
        if os.path.exists(video_path):
            os.remove(video_path)
            
        return jsonify({
            'success': True,
            'message': '视频处理成功',
            'output_directory': output_dir
        })
            
    except Exception as e:
        return jsonify({'error': f'处理视频时出错: {str(e)}'}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)
